#include <iostream>
#include <string>
#include <ctime>
#include <cstring>
#include <unistd.h>
#include <csignal>

// =============== Windows ==============
#if defined(_WIN32) || defined(_WIN64)
#include <cstdint>
#include <ws2tcpip.h>
#include <winsock2.h>
#include <windows.h>
#pragma comment(lib, "ws2_32.lib")
#define close(sock) closesocket(sock)
typedef SOCKET socket_t;
#define INVALID_SOCKET_VALUE INVALID_SOCKET
using ushort = unsigned short;
// =========  Unix/Linux/Macos  =========
#else
#include <arpa/inet.h>
#include <sys/socket.h>
#include <unistd.h>
#include <netdb.h>
#include <sys/select.h>
#define INVALID_SOCKET_VALUE -1
typedef int socket_t;
#endif

#define BUFFER_SIZE 1024

volatile std::sig_atomic_t running = 1; // 标记程序是否运行

// 优雅退出（仅服务器端 Ctrl+C 触发）
void signalHandler(int signum) {
	std::cout << "\nsignal " << signum << ", shutting down server..." << std::endl;
	running = 0; // 通知主循环退出（但客户端断开不会触发）
	exit(0);     // 强制终止进程（服务器端 Ctrl+C 时）
}

// 编译：
// Linux:   g++ -std=c++11 -o tcp_server_linux.exe tcp_server.cpp
// Windows: g++ -std=c++11 -o tcp_server_windows.exe tcp_server.cpp -lws2_32
int main(int argc, char *argv[]) {
	signal(SIGINT, signalHandler);  // Ctrl+C 触发服务器退出
	signal(SIGTERM, signalHandler); // kill 命令触发服务器退出

	if (argc != 2) {
		std::cerr << "Usage: " << argv[0] << " <port>" << std::endl;
		return 1;
	}

	int                port = std::stoi(argv[1]);
	socket_t           serverFd;
	struct sockaddr_in serverAddr;

	// 初始化 Winsock（仅 Windows）
#if defined(_WIN32) || defined(_WIN64)
	WSADATA wsaData;
	if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
		std::cerr << "WSAStartup failed!" << std::endl;
		return 1;
	}
#endif

	// 创建 TCP socket
	serverFd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (serverFd == INVALID_SOCKET_VALUE) {
		std::cerr << "TCP socket creation failed!" << std::endl;
		return 1;
	}

	// 设置 SO_REUSEADDR 选项
	int opt = 1;
	if (setsockopt(serverFd, SOL_SOCKET, SO_REUSEADDR, (const char *) &opt, sizeof(opt)) < 0) {
		std::cerr << "Setsockopt failed!" << std::endl;
		close(serverFd);
		return 1;
	}

	// 绑定地址和端口
	serverAddr.sin_family      = AF_INET;
	serverAddr.sin_port        = htons(port);
	serverAddr.sin_addr.s_addr = INADDR_ANY;
	if (::bind(serverFd, (struct sockaddr *) &serverAddr, sizeof(serverAddr)) < 0) {
		std::cerr << "TCP bind failed!" << std::endl;
		close(serverFd);
		return 1;
	}

	// 开始监听
	if (listen(serverFd, SOMAXCONN) < 0) {
		std::cerr << "TCP listen failed!" << std::endl;
		close(serverFd);
		return 1;
	}
	std::cout << "Listening on port " << port << "..." << std::endl;

	while (running) {
		struct sockaddr_in clientAddr;
		socklen_t          clientLen = sizeof(clientAddr);

		// 接受客户端连接
		// std::cout << "\n[Server] Waiting for new client..." << std::endl;
		socket_t clientFd = accept(serverFd, (struct sockaddr *) &clientAddr, &clientLen);
		if (clientFd == INVALID_SOCKET_VALUE) {
			if (!running) break; // 如果是信号触发的，退出循环
			std::cerr << "[Error] Accept failed, retrying..." << std::endl;
			continue;
		}
		std::cout << "Connected: " << inet_ntoa(clientAddr.sin_addr) << ":" << ntohs(clientAddr.sin_port) << std::endl;

		// 客户端通信循环
		char buffer[BUFFER_SIZE];
		bool clientActive = true;
		while (clientActive && running) {
			memset(buffer, 0, BUFFER_SIZE);
			ssize_t received = recv(clientFd, buffer, BUFFER_SIZE - 1, 0);
			if (received == 0) {
				std::cout << "[Error] Disconnected (EOF)." << std::endl;
				clientActive = false;
				break;
			} else if (received < 0) {
				std::cerr << "[Error] Receive failed!" << std::endl;
				clientActive = false;
				break;
			}

			buffer[received] = '\0';
			std::cout << "Received: " << buffer << std::endl;

			// 发送响应
			std::string response  = "Reply: " + std::string(buffer);
			ssize_t     sentBytes = send(clientFd, response.c_str(), response.length(), 0);
			if (sentBytes < 0) {
				std::cerr << "[Error] Send failed!" << std::endl;
				clientActive = false;
			}
		}

		// 关闭客户端连接（不影响服务器）
		close(clientFd);
		std::cout << "Client connection closed.\n" << std::endl;
	}

	// 清理服务器资源
	close(serverFd);
	std::cout << "Shutdown complete." << std::endl;
#if defined(_WIN32) || defined(_WIN64)
	WSACleanup();
#endif
	return 0;
}
